package com.study.spark.scala.ml.regression

import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.regression.LinearRegression
import org.apache.spark.sql.SparkSession

import scala.util.Random

/**
  * 线性回归算法
 *  两个变量之间存在一次方函数关系
  * @author stephen
  * @create 2019-04-07 13:05
  * @since 1.0.0
  */
object LinearRegressionDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("Linear Regression Demo")
      .master("local[2]")
      .getOrCreate()
    val file = spark.read
      .format("csv")
      .option("sep", ";")
      .option("header", "true")
      .load("/Users/stephen/Documents/03code/java-demo/bigdata-study/study-spark/src/main/resource/regression/house.csv")
    // 隐式转换
    import spark.implicits._
    val random = new Random()
    // 通过面积预测价格
    val data = file.select("square", "price")
      .map(row => (row.getAs[String](0).toDouble, row.getAs[String](1).toDouble, random.nextDouble()))
      .toDF("square", "price","random")
      // 已有的数据是排序的，需要打乱顺序，否则训练效果很差
      .sort("random")

    // 封装参数
    val assembler = new VectorAssembler()
      .setInputCols(Array("square"))
      .setOutputCol("features")
    val dataset = assembler.transform(data)

    // 数据集拆分
    val Array(train,test) = dataset.randomSplit(Array(0.8,0.2),1234L)
    // 设定参数
    val regression = new LinearRegression()
      .setMaxIter(10)
      .setRegParam(0.3)
      .setElasticNetParam(0.8)
    // 训练模型
    val model = regression.setLabelCol("price").setFeaturesCol("features").fit(train)
    // 测试集预测结果
    val result = model.transform(test)
    result.show()

  }

}
